import os
import re
import shutil
import numpy as np
import cv2
import imageio
from matplotlib import cm
from matplotlib.colors import LinearSegmentedColormap
import json

import torch

from utils.obj import write_obj
def save_img_sequence(filename, img_dir, matcher, save_format='mp4', fps=30):
    assert save_format in ['gif', 'mp4']
    if not filename.endswith(save_format):
        filename += f".{save_format}"
    matcher = re.compile(matcher)
    imgs = []
    for f in os.listdir(img_dir):
        if matcher.search(f):
            imgs.append(f)
    imgs = sorted(imgs, key=lambda f: int(matcher.search(f).groups()[0]))
    imgs = [cv2.imread(os.path.join(img_dir, f)) for f in imgs]
    
    if save_format == 'gif':
        imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
        imageio.mimsave(img_dir, imgs, fps=fps, palettesize=256)
    elif save_format == 'mp4':
        H, W, _ = imgs[0].shape
        writer = cv2.VideoWriter(filename, cv2.VideoWriter_fourcc(*'mp4v'), 30, (W, H), True)
        for img in imgs:
            writer.write(img)
        writer.release()
def main():
    save_img_sequence(filename = 'block6.mp4', img_dir = './final_test', matcher='(\d+)\.png')

if __name__ == '__main__':
    main()